package com.geek.dubbo.study;

import io.github.resilience4j.bulkhead.Bulkhead;
import io.github.resilience4j.bulkhead.BulkheadConfig;
import io.github.resilience4j.bulkhead.BulkheadRegistry;
import org.apache.dubbo.common.URL;
import org.apache.dubbo.common.constants.CommonConstants;
import org.apache.dubbo.common.extension.Activate;
import org.apache.dubbo.rpc.*;
import org.apache.dubbo.rpc.filter.ExecuteLimitFilter;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.function.Supplier;

import static org.apache.dubbo.rpc.Constants.EXECUTES_KEY;

/**
 * @Author: ysq
 * @Date: 2021/9/17
 * @Description: 通过 Dubbo Filter 机制实现 Dubbo 服务提供方限流，利用resilience4j 来整合实现
 * <p>
 * 服务提供者调用限制，使用Resilience4j的{@link Bulkhead}来实现。
 * 过滤器的代码参考了{@link ExecuteLimitFilter}的具体逻辑。
 * 该实现直接继承了{@link ExecuteLimitFilter}，原因是为了复用里面使用RpcStatus进行接口调用统计的逻辑
 */
@Activate(group = CommonConstants.PROVIDER, value = Constants.EXECUTES_KEY)
public class Resilience4jExecuteLimitFilter extends ExecuteLimitFilter {

  private static final ConcurrentMap<String, ConcurrentMap<String, Bulkhead>> Bulkhead_Instances = new ConcurrentHashMap<>();

  private static final String EXECUTE_LIMIT_FILTER_START_TIME = "execute_limit_filter_start_time";

  private static final Bulkhead getBulkhead(URL url, String methodName, int maxConcurrentCalls) {
    String uri = url.toIdentityString();
    ConcurrentMap<String, Bulkhead> map = Bulkhead_Instances.computeIfAbsent(uri, k -> new ConcurrentHashMap<>());
    return map.computeIfAbsent(methodName, k -> {
      BulkheadConfig config = BulkheadConfig.custom()
              .maxConcurrentCalls(maxConcurrentCalls)
              .build();

      Bulkhead bulkhead = BulkheadRegistry.of(config)
              .bulkhead(String.format("provider[%s]method[%s]-bulkhead", url.toIdentityString(), methodName));

      bulkhead.getEventPublisher().onCallRejected(e -> {
        // TODO 实际使用时，需要使用日志框架用debug级别去打印
        System.out.println(bulkhead.getName() + " on call rejected");
        throw new RpcException(RpcException.LIMIT_EXCEEDED_EXCEPTION,
                "Failed to invoke method " + methodName + " in provider " +
                        url + ", cause: The service using threads greater than <dubbo:service executes=\"" + maxConcurrentCalls +
                        "\" /> limited.");
      });
      bulkhead.getEventPublisher().onCallPermitted(e -> {
        // TODO 实际使用时，需要使用日志框架用debug级别去打印
        System.out.println(bulkhead.getName() + " on call perimitted");
      });
      return bulkhead;
    });
  }


  @Override
  public Result invoke(Invoker<?> invoker, Invocation invocation) throws RpcException {
    URL url = invoker.getUrl();
    String methodName = invocation.getMethodName();
    int max = url.getMethodParameter(methodName, EXECUTES_KEY, 0);

    Supplier<Result> supplier = () -> {
      try {
        return invoker.invoke(invocation);
      } catch (Throwable t) {
        if (t instanceof RuntimeException) {
          throw (RuntimeException) t;
        } else {
          throw new RpcException("unexpected exception when Resilience4jExecuteLimitFilter", t);
        }
      }
    };

    // 这里即使max大于0，仍然调用RpcStatus.beginCount，传入max为0，表示不走原来dubbo内部的限制逻辑，限制的的逻辑由Resilience4j的Bulkhead实现
    RpcStatus.beginCount(url, methodName, 0);
    // 此处是为了RpcStatus进行接口调用统计
    invocation.put(EXECUTE_LIMIT_FILTER_START_TIME, System.currentTimeMillis());

    if (max <= 0) {
      return supplier.get();
    } else {
      Bulkhead bulkhead = getBulkhead(url, methodName, max);
      return bulkhead.executeSupplier(supplier);
    }
  }
}
